import os
from random import choice, uniform

import cv2
import numpy as np
from matplotlib import pyplot as plt
from noise import pnoise2


def cv_show(img, title=None, save=False):
    if save:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        plt.imsave(title + '.png', img, dpi=600)
    cv2.imshow(title, img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


class MaskImage:
    def __init__(self, path, octaves=2, persistence=0.5, lacunarity=2.0):
        self.octaves = octaves
        self.persistence = persistence
        self.lacunarity = lacunarity

        self.path = path
        self.img = cv2.imread(self.path)
        self.shape = self.img.shape[0:2]

    def _binary_mask(self):
        gray = cv2.cvtColor(self.img, cv2.COLOR_BGR2GRAY)
        _, binary = cv2.threshold(gray, 100, 255, cv2.THRESH_OTSU | cv2.THRESH_BINARY)
        return binary

    def _noise_generate(self):
        temp = np.zeros(self.shape)
        for i in range(self.shape[0]):
            for j in range(self.shape[1]):
                temp[i][j] = pnoise2(i / 100, j / 100,
                                     octaves=self.octaves,
                                     persistence=self.persistence,
                                     lacunarity=self.lacunarity,
                                     repeatx=self.shape[0],
                                     repeaty=self.shape[1])
        temp *= 100
        temp = temp.astype(dtype=np.uint8)
        _, temp = cv2.threshold(temp, 120, 255, cv2.THRESH_BINARY)
        return temp

    def _mask_generate(self):
        m_i = self._binary_mask()
        m_p = self._noise_generate()
        m_m = cv2.bitwise_and(m_i, m_p)
        return m_m

    def process(self):
        return self._mask_generate()


def noise_foreground_generate(ori, mask, texture, trans_range=(0.15, 1)):
    texture_image_path = choice(os.listdir(texture))
    random_num = uniform(1, 10)
    if random_num <= 5:
        texture_image = texture_image_path
        i_n = cv2.imread(texture + texture_image)
        i_n = cv2.resize(i_n, ori.shape[:2])
    else:
        i_n = ori

    factor = uniform(trans_range[0], trans_range[1])
    i_n_r_1, i_n_r_2 = cv2.bitwise_and(i_n, i_n, mask=mask), cv2.bitwise_and(ori, ori, mask=mask)
    i_n_r = factor * i_n_r_1 + (1 - factor) * i_n_r_2
    return i_n_r


def simulated_generate(mask, ori, noisy):
    mask = np.where(mask == 255, 0, 255)
    i_r = cv2.bitwise_and(ori, ori, mask=mask)
    cv_show(i_r, 'i_r')
    res = i_r + noisy
    return res
